import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier  # 随机森林算法（分类器）
import matplotlib.pyplot as plt
plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15
def Feature_ranking():
    df = pd.read_csv("../data/train.csv")
    feature_data = df.copy()
    # print(df.head())
    # print(df.info())

    # 1.划分数据集
    feature_columns = feature_data.columns[1:]
    x = feature_data[feature_columns]
    y = feature_data['Attrition']

    x = pd.get_dummies(x, drop_first=True)
    x = x.drop(columns=['EmployeeNumber','StandardHours'])
    feature_columns = x.columns

    # 2.划分训练集和测试集
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=18)

    # 3.训练随机森林模型
    model = RandomForestClassifier(
        n_estimators=100,  # 决策树数量（默认100，可调整）
        random_state=18)  # 固定随机种子，结果可复现
    model.fit(x_train, y_train)


    # 4.提取特征重要性得分
    feature_importance = pd.DataFrame({
        "特征": feature_columns,
        "重要性得分": model.feature_importances_  # 模型内置的特征重要性属性
    })
    # feature_importances = np.sort(model.feature_importances_)
    # 5.按重要性得分降序排序
    feature_importance = feature_importance.sort_values(by="重要性得分", ascending=True)

    print("特征重要性排名（降序）：")
    print(feature_importance)
    feature_columns = feature_importance.iloc[:,0]
    feature_importance = feature_importance.iloc[:,-1]


    # 6.绘图分析
    plt.figure(figsize = (30,25))
    plt.barh(feature_columns,feature_importance, color='blue')
    # 绘制网格
    plt.grid(linestyle='--', color='lightgray', alpha=0.5)
    # 绘制标题，x轴 y轴标签
    plt.title("特征重要性", fontsize=15)
    plt.xlabel("重要性", fontsize=15)
    plt.ylabel("类别", fontsize=15)
    plt.show()

if __name__ == '__main__':
    Feature_ranking()


